/* *
 * @file    wakeword.c
 * @author  Azolla (1228449928@qq.com)
 * @brief   audio wake word.
 * @version 0.1
 * @date    2024-10-22
 * 
 * @copyright Copyright (c) 2024
 * */

#include <string.h>
#include "dirent.h" 

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/queue.h" 
#include "freertos/semphr.h" 
#include "freertos/event_groups.h"
#include "esp_log.h"
#include "sdkconfig.h"

#include "esp_wn_iface.h"
#include "esp_wn_models.h"
#include "esp_afe_sr_models.h"
#include "esp_mn_iface.h"
#include "esp_mn_models.h"
#include "model_path.h"

#include "audio_tts.h"
#include "hal_adc.h"
#include "sparkai.h"

#include "wakeword.h"

#define TAG  "wakeword"

#ifdef TAG
#define LOGI(format, ...)    ESP_LOGI(TAG, format, ##__VA_ARGS__)
#define LOGW(format, ...)    ESP_LOGW(TAG, format, ##__VA_ARGS__)
#define LOGE(format, ...)    ESP_LOGE(TAG, format, ##__VA_ARGS__)
#else
#define LOGI(format, ...)
#define LOGW(format, ...)
#define LOGE(format, ...)
#endif

#include "esp_heap_caps.h"
// #define MEMORY_CAPS_TYPE        (MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT)  // PSRAM
#define MEMORY_CAPS_TYPE         (MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT)     // RAM
#define memory_malloc(size)      heap_caps_malloc(size, MEMORY_CAPS_TYPE)
#define memory_calloc(n, size)   heap_caps_calloc(n, size, MEMORY_CAPS_TYPE)
#define memory_realloc(size)     heap_caps_realloc(ptr, size, MEMORY_CAPS_TYPE)
#define memory_free(ptr)         heap_caps_free(ptr)
 
static esp_afe_sr_iface_t *afe_handle = NULL;
static esp_afe_sr_data_t  *afe_data = NULL;

static const uint8_t WAKE_DETECT_EVENT = BIT0;
static EventGroupHandle_t xEvent = NULL;

bool wakeword_detect_event(void)
{
    if (xEvent == NULL) return false;
    EventBits_t uxBits = xEventGroupWaitBits(xEvent, WAKE_DETECT_EVENT, pdFALSE, pdFALSE, 0); // portMAX_DELAY
    return (uxBits & WAKE_DETECT_EVENT);
}


void feed_Task(void *arg)
{
    esp_afe_sr_data_t *afe_data = arg;
    int audio_chunksize = afe_handle->get_feed_chunksize(afe_data);
    int feed_channel = afe_handle->get_channel_num(afe_data);
    uint16_t read_size = audio_chunksize * sizeof(int16_t) * feed_channel;
    LOGI("audio_size = %d", read_size);
    adc_val_t val;
    //val.buff = memory_malloc(read_size);
    adc_read_start();
    while (1) {
        if (adc_read_buff(&val, 100)) {
            bool detect = wakeword_detect_event();
            audioai_req_handler(detect, val.buff, val.size);
            if (detect == false) {
                afe_handle->feed(afe_data, (int16_t *)val.buff);
            }
        }
        vTaskDelay(10);  // 40ms = 10ms * 1208/320
    }
    vTaskDelete(NULL);
}

void detect_Task(void *arg)
{
    //bool wakenet_is_detect = false;
    esp_afe_sr_data_t *afe_data = arg;
    int afe_chunksize = afe_handle->get_fetch_chunksize(afe_data);
    int16_t *buff = memory_malloc(afe_chunksize * sizeof(int16_t));
    assert(buff);
    LOGI("------------detect start------------");
    while (1) {

        afe_fetch_result_t *res = afe_handle->fetch(afe_data);
        if (!res || res->ret_value == ESP_FAIL) {
            LOGE("------------fetch error!------------");
            continue;
        }

        if (res->wakeup_state == WAKENET_DETECTED) {
            LOGI("-----------LISTENING-----------");
            //LOGI("model index:%d, word index:%d", res->wakenet_model_index, res->wake_word_index);
            afe_handle->disable_wakenet(afe_data);
            xEventGroupSetBits(xEvent, WAKE_DETECT_EVENT);
            audio_tts_player("请说"); 
            audioai_speech_over();  // 等待识别完成！
            afe_handle->enable_wakenet(afe_data);
            xEventGroupClearBits(xEvent, WAKE_DETECT_EVENT);
            LOGI("vad silence...");
        }
        // LOGI("vad_state = %d", res->vad_state);
        // res->vad_state == AFE_VAD_SILENCE noise or silence
        // vTaskDelay(1000);
    }
    vTaskDelete(NULL);
}

void wakeword_init(void)
{
    srmodel_list_t *models = esp_srmodel_init("model");
    char *wn_name = NULL;
    char *wn_name_2 = NULL;
    if (models != NULL) {
        for (int i = 0; i<models->num; i++) {
            if (strstr(models->model_name[i], ESP_WN_PREFIX) != NULL) {
                if (wn_name == NULL) {
                    wn_name = models->model_name[i];
                    LOGI("The first wakenet model: %s", wn_name);
                } else if (wn_name_2 == NULL) {
                    wn_name_2 = models->model_name[i];
                    LOGI("The second wakenet model: %s", wn_name_2);
                }
            }
        }
    } else {
        LOGW("Please enable wakenet model and select wake word by menuconfig!");
        return ;
    }

    afe_handle = (esp_afe_sr_iface_t *)&ESP_AFE_SR_HANDLE;
    afe_config_t afe_config = AFE_CONFIG_DEFAULT();
    afe_config.memory_alloc_mode = AFE_MEMORY_ALLOC_MORE_PSRAM;
    afe_config.wakenet_init = true;
    afe_config.wakenet_model_name = wn_name;
    afe_config.wakenet_model_name_2 = wn_name_2;
    afe_config.voice_communication_init = false;

    afe_config.pcm_config.total_ch_num = 1;
    afe_config.pcm_config.mic_num = 1;
    afe_config.pcm_config.ref_num = 0;
    afe_config.aec_init = false;
    afe_config.wakenet_mode = DET_MODE_90;
    afe_config.se_init = false;
 
    afe_data = afe_handle->create_from_config(&afe_config);

    xEvent = xEventGroupCreate();
    xTaskCreatePinnedToCore(&feed_Task, "feed", 8 * 1024, (void*)afe_data, 10, NULL, APP_CPU_NUM);
    xTaskCreatePinnedToCore(&detect_Task, "detect", 4 * 1024, (void*)afe_data, 9, NULL, PRO_CPU_NUM);
}